// COMPI Collectives Benchmarks
// Fixed vs Non-fixed, Small vs Large, Balanced vs Unbalanced

#include "compi/collectives.h"
#include "compi/compi.h"

#include <algorithm>
#include <cstddef>
#include <cstdio>
#include <cstdlib>
#include <string>
#include <vector>

using namespace compi;
using namespace compi::collectives;

struct BenchConfig
{
    int iterations = 10;
    int small_count = 16;    // elements per rank (balanced)
    int large_count = 8192;  // elements per rank (balanced)
    int small_str_len = 16;  // average string length
    int large_str_len = 256; // average string length
    int imbalance_delta = 7; // introduce count skew (unbalanced)
};

static std::vector<int>
make_fixed_values(int count, int rank)
{
    std::vector<int> v(count);
    for (int i = 0; i < count; ++i)
        v[i] = rank * 100000 + i;
    return v;
}

static std::string
make_text(int len, int seed)
{
    std::string s;
    s.resize(static_cast<std::size_t>(len));
    for (int i = 0; i < len; ++i)
        s[static_cast<std::size_t>(i)] = static_cast<char>('a' + ((seed + i) % 26));
    return s;
}

static std::vector<std::string>
make_strings(int count, int avg_len, bool jitter, int rank)
{
    std::vector<std::string> v;
    v.reserve(static_cast<std::size_t>(count));
    for (int i = 0; i < count; ++i)
    {
        int len = avg_len;
        if (jitter)
        {
            // mild variability to simulate non-fixed payloads
            len = std::max(1, avg_len + ((rank + i) % 5) - 2);
        }
        v.emplace_back(make_text(len, rank * 1000 + i));
    }
    return v;
}

template <typename Fn>
static double
measure(Context & ctx, int iterations, Fn && op)
{
    // synchronize before timing
    MPI_Barrier(ctx.comm());
    double t0 = MPI_Wtime();
    for (int i = 0; i < iterations; ++i)
    {
        op();
    }
    MPI_Barrier(ctx.comm());
    double t1 = MPI_Wtime();
    return (t1 - t0) / static_cast<double>(iterations);
}

static void
print_stats(Context & ctx, const char * name, double avg_time)
{
    // gather per-rank times to root for min/max/avg reporting
    auto times = gather<double>(ctx, /*root=*/0, avg_time);
    if (ctx.rank() == 0)
    {
        double sum = 0.0, mn = 1e100, mx = -1e100;
        for (double t : times)
        {
            sum += t;
            mn = std::min(mn, t);
            mx = std::max(mx, t);
        }
        double avg = sum / static_cast<double>(times.size());
        std::printf("[bench] %-28s | avg: %.6f s | min: %.6f | max: %.6f | ranks: %d\n",
                    name,
                    avg,
                    mn,
                    mx,
                    ctx.size());
        std::fflush(stdout);
    }
}

template <typename Fn>
static void
run_scenario(Context & ctx, const char * name, int iterations, Fn && op)
{
    // Announce start on root to help pinpoint crashes
    if (ctx.rank() == 0)
    {
        std::printf("[bench] starting: %s\n", name);
        std::fflush(stdout);
    }
    // Ensure ranks are aligned before running
    MPI_Barrier(ctx.comm());
    try
    {
        double t = measure(ctx, iterations, std::forward<Fn>(op));
        print_stats(ctx, name, t);
    }
    catch (const std::exception & ex)
    {
        std::fprintf(
            stderr, "[bench][rank %d] scenario '%s' threw: %s\n", ctx.rank(), name, ex.what());
        std::fflush(stderr);
    }
    catch (...)
    {
        std::fprintf(stderr, "[bench][rank %d] scenario '%s' threw: unknown\n", ctx.rank(), name);
        std::fflush(stderr);
    }
    // Re-align to avoid cascading failures
    MPI_Barrier(ctx.comm());
}

int
main(int argc, char ** argv)
{
    Environment env;
    Context & ctx = Context::for_comm(MPI_COMM_WORLD, env);
    BenchConfig cfg;

    if (argc > 1)
    {
        // optional: iterations override
        cfg.iterations = std::atoi(argv[1]);
        if (cfg.iterations <= 0)
            cfg.iterations = 10;
    }

    int rank = ctx.rank();
    int size = ctx.size();

    // Scenario A: allgather fixed single value (int)
    {
        int my = rank * 123 + 7;
        run_scenario(ctx,
                     "allgather<int> single",
                     cfg.iterations,
                     [&]()
                     {
                         auto out = allgather<int>(ctx, my);
                         (void)out;
                     });
    }

    // Scenario B1: allgatherv<int> balanced small
    {
        auto local = make_fixed_values(cfg.small_count, rank);
        run_scenario(ctx,
                     "allgatherv<int> small balanced",
                     cfg.iterations,
                     [&]()
                     {
                         auto out = allgatherv<int>(ctx, local);
                         (void)out;
                     });
    }

    // Scenario B2: allgatherv<int> balanced large
    {
        auto local = make_fixed_values(cfg.large_count, rank);
        run_scenario(ctx,
                     "allgatherv<int> large balanced",
                     cfg.iterations,
                     [&]()
                     {
                         auto out = allgatherv<int>(ctx, local);
                         (void)out;
                     });
    }

    // Scenario B3: allgatherv<int> unbalanced counts
    {
        int count = cfg.small_count + (rank % cfg.imbalance_delta);
        auto local = make_fixed_values(count, rank);
        run_scenario(ctx,
                     "allgatherv<int> small unbalanced",
                     cfg.iterations,
                     [&]()
                     {
                         auto out = allgatherv<int>(ctx, local);
                         (void)out;
                     });
    }

    // Scenario C1: allgather_packed_range<string> small balanced
    {
        auto local = make_strings(cfg.small_count, cfg.small_str_len, /*jitter=*/true, rank);
        run_scenario(ctx,
                     "packed-range<string> small balanced",
                     cfg.iterations,
                     [&]()
                     {
                         auto out = allgather_packed_range<std::string>(ctx, local);
                         (void)out;
                     });
    }

    // Scenario C2: allgather_packed_range<string> large balanced
    {
        auto local = make_strings(cfg.large_count, cfg.large_str_len, /*jitter=*/true, rank);
        run_scenario(ctx,
                     "packed-range<string> large balanced",
                     cfg.iterations,
                     [&]()
                     {
                         auto out = allgather_packed_range<std::string>(ctx, local);
                         (void)out;
                     });
    }

    // Scenario D1: allgatherv<string> small balanced (typed *v path for non-fixed)
    {
        auto local = make_strings(cfg.small_count, cfg.small_str_len, /*jitter=*/true, rank);
        run_scenario(ctx,
                     "allgatherv<string> small balanced",
                     cfg.iterations,
                     [&]()
                     {
                         auto out = allgatherv<std::string>(ctx, local);
                         (void)out;
                     });
    }

    // Scenario D2: allgatherv<string> large balanced
    {
        auto local = make_strings(cfg.large_count, cfg.large_str_len, /*jitter=*/true, rank);
        run_scenario(ctx,
                     "allgatherv<string> large balanced",
                     cfg.iterations,
                     [&]()
                     {
                         auto out = allgatherv<std::string>(ctx, local);
                         (void)out;
                     });
    }

    // Scenario D3: allgatherv<string> small unbalanced
    {
        int count = cfg.small_count + (rank % cfg.imbalance_delta);
        auto local = make_strings(count, cfg.small_str_len, /*jitter=*/true, rank);
        run_scenario(ctx,
                     "allgatherv<string> small unbalanced",
                     cfg.iterations,
                     [&]()
                     {
                         auto out = allgatherv<std::string>(ctx, local);
                         (void)out;
                     });
    }

    // Scenario E1: alltoall<int> Default (fixed)
    {
        std::vector<int> send_values(static_cast<std::size_t>(size));
        for (int i = 0; i < size; ++i)
            send_values[static_cast<std::size_t>(i)] = rank * 100 + i;
        run_scenario(ctx,
                     "alltoall<int> Default",
                     cfg.iterations,
                     [&]()
                     {
                         auto r =
                             alltoall<int>(ctx, send_values, SyncType::Default, SendMode::Default);
                         (void)r;
                     });
    }

    // Scenario E2: alltoall<int> NBX + Synchronous (fixed)
    {
        std::vector<int> send_values(static_cast<std::size_t>(size));
        for (int i = 0; i < size; ++i)
            send_values[static_cast<std::size_t>(i)] = rank * 100 + i;
        run_scenario(ctx,
                     "alltoall<int> NBX+Sync",
                     cfg.iterations,
                     [&]()
                     {
                         auto r =
                             alltoall<int>(ctx, send_values, SyncType::NBX, SendMode::Synchronous);
                         (void)r;
                     });
    }

    // Scenario E3: alltoall<string> Default (non-fixed TLV two-frame)
    {
        std::vector<std::string> send_values(static_cast<std::size_t>(size));
        for (int i = 0; i < size; ++i)
        {
            send_values[static_cast<std::size_t>(i)] =
                std::string("msg_") + std::to_string(rank) + "_to_" + std::to_string(i);
        }
        run_scenario(ctx,
                     "alltoall<string> Default",
                     cfg.iterations,
                     [&]()
                     {
                         auto r = alltoall<std::string>(
                             ctx, send_values, SyncType::Default, SendMode::Default);
                         (void)r;
                     });
    }

    // Scenario E4: alltoall<string> NBX (non-fixed, two-phase)
    {
        std::vector<std::string> send_values(static_cast<std::size_t>(size));
        for (int i = 0; i < size; ++i)
        {
            // vary payload to simulate non-fixed sizes
            int len = cfg.small_str_len + ((rank + i) % 5);
            send_values[static_cast<std::size_t>(i)] = make_text(len, rank * 1000 + i);
        }
        run_scenario(ctx,
                     "alltoall<string> NBX",
                     cfg.iterations,
                     [&]()
                     {
                         auto r = alltoall<std::string>(
                             ctx, send_values, SyncType::NBX, SendMode::Default);
                         (void)r;
                     });
    }

    if (rank == 0)
    {
        std::printf("\n[bench] Done. You can tweak iterations via argv[1].\n");
        std::fflush(stdout);
    }

    return 0;
}